import type { BaseEmbeddings } from '@cherrystudio/embedjs-interfaces'
import { KnowledgeBaseParams } from '@types'

import EmbeddingsFactory from './EmbeddingsFactory'

/**
 * 文本嵌入处理类
 * 封装了文本向量化的核心操作，通过工厂模式创建具体的嵌入实现
 */
export default class Embeddings {
  // 具体的嵌入实现实例
  private sdk: BaseEmbeddings

  /**
   * 构造函数
   * @param params 知识库参数配置
   *   - model: 使用的模型名称
   *   - apiKey: API密钥
   *   - apiVersion: API版本
   *   - baseURL: 基础URL
   *   - dimensions: 向量维度
   */
  constructor({ model, apiKey, apiVersion, baseURL, dimensions }: KnowledgeBaseParams) {
    // 通过工厂创建具体的嵌入实现
    this.sdk = EmbeddingsFactory.create({ model, apiKey, apiVersion, baseURL, dimensions } as KnowledgeBaseParams)
  }

  /** 初始化嵌入服务 */
  public async init(): Promise<void> {
    return this.sdk.init()
  }

  /** 获取向量维度 */
  public async getDimensions(): Promise<number> {
    return this.sdk.getDimensions()
  }

  /**
   * 批量文本向量化
   * @param texts 文本数组
   * @returns 向量数组(二维)
   */
  public async embedDocuments(texts: string[]): Promise<number[][]> {
    return this.sdk.embedDocuments(texts)
  }

  /**
   * 单条文本向量化
   * @param text 输入文本
   * @returns 向量数组(一维)
   */
  public async embedQuery(text: string): Promise<number[]> {
    return this.sdk.embedQuery(text)
  }
}
